/*
 *  CHEEZBURGER.cpp
 *  TubeTruck
 *
 *  Created by Aeturnum on 6/6/08.
 *  Copyright 2008 __MyCompanyName__. All rights reserved.
 *
 */

#include "CHEEZBURGER.h"

using namespace std;
using namespace TTP;

/*
std::vector<uint32_t> mIdHashList;
std::vector<uint32_t> mAckNumbers;
std::vector<AppDataSegment> mData;
*/

CHEEZBURGER::CHEEZBURGER() : Packet(Packet_CHEEZBURGER), dataIndex(-1)
{}
CHEEZBURGER::CHEEZBURGER(id_hash idHash) : Packet(idHash, Packet_CHEEZBURGER)
{}
CHEEZBURGER::CHEEZBURGER(uint32_t idHash,  vector<byte>& sig, vector<uint32_t>& idHashList, vector<uint32_t>& ackNumbers) :
Packet(idHash, Packet_CHEEZBURGER), Signable(sig), mIdHashList(idHashList), mAckNumbers(ackNumbers)
{}
CHEEZBURGER::~CHEEZBURGER() {}

//accessors
vector<id_hash> CHEEZBURGER::getHashList() const {
   return mIdHashList;
}

vector<ack_num> CHEEZBURGER::getAckNumbers() const {
   return mAckNumbers;
}

ack_num CHEEZBURGER::getAckNumber(id_hash idHash) const {
   
   for(uint32_t i = 0; i < mIdHashList.size(); i++) {
      if(mIdHashList[i] == idHash)
         return mAckNumbers[i];
   }
   
   throw msg_exception("CHEEZBURGER: Requested Ack number for non-existant idHash.");
}

vector< vector<byte> > CHEEZBURGER::getDataVectors() const {
   vector< vector<byte> > dataVectors;
   map<ack_num, AppDataSegment>::const_iterator itr = mData.begin();
   
   while(itr != mData.end()) {
      dataVectors.push_back(itr->second.getData());
      ++itr;
   }
   
   return dataVectors;
}

vector<ack_num> CHEEZBURGER::getDataAckNumbers() const {
   vector<ack_num> retVal;
   std::map<ack_num, AppDataSegment>::const_iterator itr = mData.begin();
   while(itr != mData.end()) {
      retVal.push_back(itr->second.getAckNumber());
      assert(itr->second.getAckNumber() == itr->first);
      itr++;
   }

   return retVal;
}

vector<byte> CHEEZBURGER::getDataVector(uint32_t ackNumber) const {
   map<ack_num, AppDataSegment>::const_iterator appData = mData.find(ackNumber);
   
   if(appData == mData.end())
      throw msg_exception("CHEEZEBURGER:No data found for specified ACK\n");
   
   return appData->second.getData();
}

vector<byte> CHEEZBURGER::getBlowfishIv() {
   return mBlowfishIv;
}

// modifiers
void CHEEZBURGER::addHashAckPair(id_hash idHash, ack_num ackNumber) {
   mIdHashList.push_back(idHash);
   mAckNumbers.push_back(ackNumber);
   
   if(mIdHashList.size() != mAckNumbers.size()) {
      throw msg_exception("CHEEZEBURGER:List of Hash and Ack numbers unequal in size.\n");
   }
}

void CHEEZBURGER::setAckNumber(ack_num ackNumber) {
   if(mData.find(ackNumber) == mData.end()) {
      AppDataSegment temp;
      temp.setAckNumber(ackNumber);   
      mData[ackNumber] =  temp;
   }
   dataIndex = ackNumber;
}

void CHEEZBURGER::addDataSegment(ack_num ack, std::vector<byte>& data) {
   setAckNumber(ack);
   addDataSegment(data);
}

void CHEEZBURGER::addDataSegment(vector<byte>& data) {
   mData[dataIndex].appendData(data);
}

void CHEEZBURGER::setSourceIdHash(id_hash idHash) {
   mIdHash = idHash;
}

id_hash CHEEZBURGER::getSourceIdHash() const {
   return mIdHash;
}

void CHEEZBURGER::setBlowfishIv(vector<byte>& blowfishIv) {
   mBlowfishIv = blowfishIv;
}

// inherited
void CHEEZBURGER::fromVec(vector<byte>::const_iterator &vit, vector<byte>::const_iterator &end) {
   mEncryptedSection.clear();
   mBlowfishIv.clear();
   
   Packet::fromVec(vit, end);
   mSig.fromVec(vit, end);
   
   mBlowfishIv.insert(mBlowfishIv.end(), vit, vit + 8);
   vit = vit + 8;
   
   mEncryptedSection.insert(mEncryptedSection.end(), vit, end);
   mEncrypted = true;
}

void CHEEZBURGER::toVec(vector<byte> &vec) const {
   if(!mEncrypted) {
      throw msg_exception("ERROR - toVec::CHEEZBURGER not encrypted!");
   } else {
      if(!mSigned) {
         throw msg_exception("ERROR - toVec::CHEEZBURGER not signed!");
      } else {
         if(mBlowfishIv.empty()) {
            throw msg_exception("ERROR: Blowfish IV not set!");
         } else {
            Packet::toVec(vec);
            mSig.toVec(vec);
            vec.insert(vec.end(), mBlowfishIv.begin(), mBlowfishIv.end());
            vec.insert(vec.end(), mEncryptedSection.begin(), mEncryptedSection.end());
         }
      }
   }
}

void CHEEZBURGER::setEncryptedData(vector<byte>& encData) {
   mEncryptedSection = encData;
   mEncrypted = true;
}

void CHEEZBURGER::setPlainText(vector<byte>& vec) {
   mData.clear();
   mAckNumbers.clear();
   mIdHashList.clear();
   
   vector<byte>::const_iterator vit = vec.begin(), end = vec.end();
   
   unsigned short numberOfHashAckPairs = 0;
   unserialize<unsigned short>(vit, end, numberOfHashAckPairs);
   
   for(unsigned short i = 0; i< numberOfHashAckPairs; i++) {
      uint32_t temp = 0;
      
      unserialize<uint32_t>(vit, end, temp);
      mIdHashList.push_back(temp);
      
      unserialize<uint32_t>(vit, end, temp);
      mAckNumbers.push_back(temp);
   }
   
   while(vit != vec.end()) {
      AppDataSegment temp;
      temp.fromVec(vit, end);
//      printf("\nCreating AppDataSegment from vector\n");
//      printf("Temp's data: Length = %d, ack# = %d", temp.getDataLength(), temp.getAckNumber());
      mData[temp.getAckNumber()] =  temp;
   }
   
   fillEncryptedSection();
}

void CHEEZBURGER::fillEncryptedSection() {
   unsigned short numberOfHashAckPairs = mIdHashList.size();
   // TODO - make sure the size listed here is corret!!
   uint32_t unencryptedSize = (numberOfHashAckPairs)*8+2;
   
//   cout << "Initial size: " << unencryptedSize << " numberOfHashAckPairs: " << numberOfHashAckPairs << "\n";
   
   FOR_EACH(it, mData) {
      unencryptedSize += it->second.getSize();
  //    cout << "Data Chunk: " << mData[i].getSize() << " Total size: "<< unencryptedSize << "\n";
   }
   
//   cout << "Filling section\n\n\n";
   
   if(mEncryptedSection.size() != unencryptedSize) {
      mEncryptedSection.clear();
      
      serialize<unsigned short>(mEncryptedSection, numberOfHashAckPairs);
//      cout << "serialize: " << mEncryptedSection.size() << "\n";
      
      for(uint32_t i = 0; i< numberOfHashAckPairs; i++) {
         uint32_t temp = mIdHashList[i];
         serialize<uint32_t>(mEncryptedSection, temp);
//         cout << "hash: " << mEncryptedSection.size() << "\n";
         
         temp = mAckNumbers[i];
         serialize<uint32_t>(mEncryptedSection, temp);
//         cout << "ack: " << mEncryptedSection.size() << "\n";
      }
      
      FOR_EACH(it, mData) {   
         it->second.toVec(mEncryptedSection);
      }
      
//      printf("unencryptedSize: %d, mEncryptedSection.size(): %d\n", unencryptedSize, (int)mEncryptedSection.size());
//      printf("numberOfHashAckPairs: %d, mData[0].size: %d\n",numberOfHashAckPairs, mData[0].getSize());
      
      assert(unencryptedSize == mEncryptedSection.size());
   }
}

vector<byte> CHEEZBURGER::getEncryptedSection() {
   if(!mEncrypted)
      fillEncryptedSection();
   
   return mEncryptedSection;
}

vector<byte> CHEEZBURGER::getUnsignedData() const {
   if(!mEncrypted) {
      throw msg_exception("ERROR:getUnsignedData: CHEEZBURGER not encrypted!");
   } else {
      vector<byte> dataToSign;
      vector<byte> emptySig;
      
      Packet::toVec(dataToSign);
      emptySig.insert(emptySig.end(), mSig.size(), 0);
      dataToSign.insert(dataToSign.end(), emptySig.begin(), emptySig.end());
      dataToSign.insert(dataToSign.end(), mEncryptedSection.begin(), mEncryptedSection.end());
      return dataToSign;
   }
}
